# 旋转目标检测与校正

import cv2
import os
import base64
from app import image_util
from flask import Blueprint, request, jsonify
from ultralytics import YOLO


# 创建蓝图
obb_bp = Blueprint('obb_bp', __name__)

# 获取当前脚本的绝对路径
script_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(script_dir)

# 加载多个 YOLO 模型并存储在字典中
#   models = {
#       'model1': YOLO('./runs/obb/train3/weights/best.pt')，
#       'model2': YOLO('./runs/obb/train3/weights/best1.pt')
#   }

# obb_000 : 通用分类模型，适配所有企业，用来检测表格位置和角度
# obb_001 : ***********************
# obb_002 : 通用分类模型，用来检测标题位置和角度
models = {
    'obb_000': YOLO(os.path.join(parent_dir, 'static', 'model', 'obb', 'obb_000.pt')),
    'obb_001': YOLO(os.path.join(parent_dir, 'static', 'model', 'obb', 'obb_001.pt')),
    'obb_002': YOLO(os.path.join(parent_dir, 'static', 'model', 'obb', 'obb_002.pt'))
}


# web api 执行推理
@obb_bp.route('/inference_image', methods=['POST', 'GET'])
def inference_image():
    """
    通过OBB 模型校正图片
    :return:  image base64格式
    """
    img_path = ''
    model_code = ''
    # 判断请求类型 获取请求参数
    if request.method == 'GET':
        # 获取 GET 请求参数
        img_path = request.args.get('imgPath')
        model_code = request.args.get('modelCode')
    elif request.method == 'POST':
        # 获取表单数据
        if request.form:
            data = request.form
            img_path = data.get('imgPath')
            model_code = data.get('modelCode')
        # 获取 JSON 数据
        elif request.json:
            data = request.json
            img_path = data.pop('imgPath')
            model_code = data.pop('modelCode')

    else:
        return "Unsupported request method"

    print(img_path)
    print(model_code)
    # 获取图片
    image = image_util.read_image(img_path)
    if model_code not in models:
        return jsonify({'error': f'Model {model_code} not found'}), 404

    # 使用指定的模型进行预测
    table_model = models[model_code]
    table_res = table_model(image)
    title_model = models['obb_002']
    title_res = title_model(image)

    # 两次校正（先转正在微调）效果会好一点
    # 校正图片--将图片方向转到正方向
    image = image_util.correct_image_to_up(image, table_res, title_res)

    # 对转正的图片进行推理
    table_res = table_model(image)
    # 对转正的图片进行校正 --对正方向图片进行角度微调
    image = image_util.correct_image(image, table_res)

    # cv2.imshow('Corrected Image', image)
    # cv2.waitKey(0)

    # 将 numpy.ndarray 图像转换为字节流
    _, buffer = cv2.imencode('.jpg', image)
    image_data = buffer.tobytes()
    return base64.b64encode(image_data).decode('utf-8')
